/* merge.c
 * This file is part of Landmerge (http://code.google.com/p/landmerge/)
 * Copyright (c) Matus Tejiscak 2008
 * Covered by GNU GPL v3 (http://www.gnu.org/licenses/gpl.txt).
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdarg.h>
#include <string.h>

#include "tiff.h"

// Symbolic indexes to various structures
#define R 0
#define G 1
#define B 2
#define P 3

/// Set of input TIFFs.
typedef struct tiff_in_t
{
	tiff_t *chan[4];
}
tiff_in_t;

/// Image statistics. (Applies on grayscale images only).
typedef struct stats_t
{
	/// \sum x_i
	double Sx;
	/// avg(x_i)
	double avg;
	/// \sum 1
	unsigned long pixel_count;

	/// \sum (x_i - avg_x)^2
	double Sdd;
	/// standard deviation
	double sd;
}
stats_t;

// Crop bounds
int crop_x = 2500;
int crop_y = 6200;
int crop_w = 3000;
int crop_h = 3000;
int tlx, tly, trx, try, blx, bly, brx, bry;

// Constants		  R    G    B    Pan
/// Gamma correction.
double gamma[4] =	{ 1.0, 0.9, 1.0, 1.4 };
/// <0..1> pixel range will be mapped to <avg - dyn_range*SD, avg + dyn_range*SD>
/// where SD is standard deviation of the samples.
double dyn_range[4] =	{ 3.0, 3.0, 3.0, 4.0 };
/// Channel subsampling. Constant, according to the Landsat docs.
int subsamp[4] =	{ 2,   2,   2,   1   };
/// 0 = leave original Y, 1 = replace by P completely
double p_override = 0.75;
/// Original Y's gamma correction.
double ygamma = 1.4;

/// Print an error message and quit.
void die(const char * const reason, ...)
{
	va_list ap;
	va_start(ap, reason);
	vfprintf(stderr, reason, ap);
	va_end(ap);
	exit(1);
}

/// An utility function to print progress.
void progress(const char * const fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	vprintf(fmt, ap);
	va_end(ap);
	fflush(stdout);
}

__attribute__ ((gnu_inline)) inline double sqr(double x)
{
	return x * x;
}

__attribute__ ((gnu_inline)) inline double apply_gamma(double x, double gamma)
{
	return  (x < 1.0e-9) ? 0 : exp(gamma * log(x));
}

#define MIN(a, b) ((a < b) ? (a) : (b))
#define MAX(a, b) ((a > b) ? (a) : (b))
#define BOUND(x) ((x < 0) ? 0 : ((x > 1) ? 1 : x))
#define QUANT(x) lround(BOUND(x) * 255)
/// Calculate one pixel's RGB color
void calculate_pixel(unsigned char *pixel, unsigned char *val, stats_t *stats)
{
	double x[4];

	if (!val[R] && !val[G] && !val[B])
	{
		pixel[R] = pixel[G] = pixel[B] = 0;
		return;
	}

	// Map [0..255] to <0.0, 1.0>
	for (int i = 0; i < 4; ++i)
		x[i] = val[i] / 255.0;

	// Normalize contrast & apply gamma
	for (int i = 0; i < 4; ++i)
	{
		x[i] = (1 + (x[i] - stats[i].avg) / (dyn_range[i]*stats[i].sd)) / 2;
		x[i] = apply_gamma(x[i], gamma[i]);
	}

	// Convert to YPbPr
	// taken from http://en.wikipedia.org/wiki/Y'CbCr
	double Y, Pb, Pr;
	Y  = + 0.299    * x[R] + 0.587    * x[G] + 0.114    * x[B];
	Pb = - 0.168736 * x[R] - 0.331264 * x[G] + 0.5      * x[B];
	Pr = + 0.5      * x[R] - 0.418688 * x[G] - 0.081312 * x[B];

	// Apply Y gamma
	Y = apply_gamma(Y, ygamma);

	// Enhance the resolution (combine the luminance channel)
	Y = p_override * x[P] + (1-p_override) * Y;

	// Convert back to RGB
	// this is the inverse of the former matrix
	x[R] = Y + 1.402 * Pr;
	x[G] = Y - 0.344 * Pb - 0.714 * Pr;
	x[B] = Y + 1.772 * Pb;

	// Write the pixel
	for (int i = 0; i < 3; ++i)
		pixel[i] = QUANT(x[i]);
}

/// Sums pixels in the given scanline. (Grayscale only).
void cb_sum_pixels(const unsigned char *begin, const unsigned char *end, stats_t *stats)
{
	unsigned line_sum = 0;

	for (const unsigned char *p = begin; p < end; ++p)
	{
		if (*p == 0) continue;

		line_sum += *p;
		++stats->pixel_count;
	}

	// Add by lines to avoid rounding errors.
	stats->Sx += line_sum / 255.0;
}

/// Sums (x_i - avg_x)^2 over the given scanline. (Grayscale only).
void cb_sum_diff_squares(const unsigned char *begin, const unsigned char *end, stats_t *stats)
{
	double line_sum = 0;

	for (const unsigned char *p = begin; p < end; ++p)
	{
		if (*p == 0) continue;

		double val = *p / 255.0;
		line_sum += sqr(val - stats->avg);
	}

	// Add by lines to avoid rounding errors.
	stats->Sdd += line_sum;
}

/// Executes a given function for each scanline.
void foreach_scanline(tiff_t *tiff, stats_t *stats, int subsamp, void (*callback)(const unsigned char *, const unsigned char *, stats_t *))
{
	const unsigned char *buf = (const unsigned char *) tiff->scanline;
	unsigned width = tiff->row_size;
	unsigned height = tiff->row_count;

	int cr_y = crop_y / subsamp, cr_h = crop_h / subsamp;
	int cr_x = crop_x / subsamp, cr_w = crop_w / subsamp;

	for (int y = cr_y; y < height && y < cr_y + cr_h; ++y)
	{
		tiff_seek(tiff, y);
		callback(buf + cr_x, buf + MIN(cr_x + cr_w, width), stats); 
		if (y % 512 == 0) progress("gathering at line %d\r", y);
	}
	printf("                                                      \r");
}

/// Gathers image statistics.
void gather_stats(tiff_t *tiff, stats_t *stats, int subsamp)
{
	printf("gathering stats...\n");

	// 1st pass -- sum && average
	stats->Sx = 0;
	stats->pixel_count = 0;
	foreach_scanline(tiff, stats, subsamp, cb_sum_pixels);
	stats->avg = stats->Sx / stats->pixel_count;

	printf("  n = %d\n", stats->pixel_count);
	printf("  Sx = %g\n", stats->Sx);
	printf("  avg = %g\n", stats->avg);

	// 2nd pass -- variance && standard deviation
	stats->Sdd = 0;
	foreach_scanline(tiff, stats, subsamp, cb_sum_diff_squares);
	stats->sd = sqrt(stats->Sdd / (stats->pixel_count - 1));

	printf("  Sdd = %g\n", stats->Sdd);
	printf("  sd = %g\n", stats->sd);
}

/// Merges all TIFFs into a single RGB file.
void merge(const tiff_in_t *tiffs, FILE *out_rgb)
{
	printf("merging...\n");
	// The output size is the panchromatic image's size
	unsigned width = tiffs->chan[P]->row_size;
	unsigned height = tiffs->chan[P]->row_count;

	unsigned outline_size = width * 3;
	unsigned char *outline = (unsigned char *) malloc(outline_size);

	// Gather stats for later use
	stats_t stats[4];
	for (int col = 0; col < 4; ++col)
	{
		gather_stats(tiffs->chan[col], stats + col, subsamp[col]);
	}

	// Initialize scanlines
	unsigned char *next[4], *current[4];
	for (int col = 0; col < 4; ++col)
	{
		unsigned size = tiffs->chan[col]->row_size;

		tiff_seek(tiffs->chan[col], crop_y / subsamp[col] - 1);

		next[col] = (unsigned char *) tiffs->chan[col]->scanline;
		current[col] = (unsigned char *) malloc(tiffs->chan[col]->row_size);
		memcpy(current[col], next[col], tiffs->chan[col]->row_size);

		tiff_next_scanline(tiffs->chan[col]);
	}

	// Merge the image
	for (int y = crop_y; y < crop_y + crop_h && y < height; ++y)
	{
		// Determine scene borders
		int bound_left, bound_right;

		// Clear the outline
		memset(outline, 0, outline_size);

		// Left border
		if (y > bly)
		{
			int Dx = brx - blx;
			int dy = y - bly;
			int Dy = bry - bly;
			bound_left = brx + dy*Dx/Dy;
		}
		else
		{
			int Dx = blx - tlx;
			int dy = y - tly;
			int Dy = bly - tly;
			bound_left = tlx + dy*Dx/Dy;
		}

		// Right border
		if (y > try)
		{
			int Dx = brx - trx;
			int Dy = bry - try;
			int dy = y - try;
			bound_right = trx + dy*Dx/Dy;
		}
		else
		{
			int Dx = trx - tlx;
			int Dy = try - tly;
			int dy = y - tly;
			bound_right = tlx + dy*Dx/Dy;
		}
		
		// Read next scanline(s) if needed
		for (int col = 0; col < 4; ++col)
		{
			if (y % subsamp[col] == 0)
			{
				tiffs->chan[col]->scanline = current[col];
				current[col] = next[col];
				next[col] = tiffs->chan[col]->scanline;
				tiff_next_scanline(tiffs->chan[col]);
			}
		}

		// Where the pixel will be written
		unsigned char *pixel = outline + 3*crop_x;
		unsigned char val[4];
		// Traverse all pixels
		for (int x = crop_x; x < crop_x + crop_w && x < width; ++x)
		{
			// Skip invalid area
			if (x < bound_left || x > bound_right)
			{
				pixel += 3;
				continue;
			}

			int x_sampled[4];
			for (int col = 0; col < 4; ++col)
			{
				// Linear interpolation: calculate the four pixels
				int sx = x / subsamp[col];
				unsigned char top_left = current[col][sx];
				unsigned char top_right = (x < width-1) ? current[col][sx+1] : top_left;
				unsigned char bottom_left = next[col][sx];
				unsigned char bottom_right = (x < width-1) ? next[col][sx+1] : bottom_left;
				int xmod = x % subsamp[col];
				int ymod = y % subsamp[col];

				// Interpolate horizontally
				unsigned char upper = (xmod*top_right + (subsamp[col]-xmod)*top_left) / subsamp[col];
				unsigned char lower = (xmod*bottom_right + (subsamp[col]-xmod)*bottom_left) / subsamp[col];

				// Interpolate vertically
				val[col] = (ymod*lower + (subsamp[col]-ymod)*upper) / subsamp[col];
			}

			// Calculate & write pixel value
			calculate_pixel(pixel, val, stats);

			// Move on...
			pixel += 3;
		}

		// Write the cropped portion of the scanline
		fwrite(outline + crop_x*3, crop_w*3, 1, out_rgb);

		// Print some progress
		if (y % 256 == 0) progress("Line %d merged.\r", y);
	}
	printf("merging done.                    \n");

	// Free memory
	free(outline);
	for (int col = 0; col < 4; ++col)
		free(current[col]);
}

/// Prints usage info & quits
void usage()
{
	printf("usage:\n");
	printf("\t./merge red.tif green.tif blue.tif panchrom.tif <crop_definition> output_file.rgb <scene_rect>\n");
	printf("crop definition:\n\tWIDTHxHEIGHT@X,Y  e.g. 800x600@1017,534\n");
	exit(1);
}

/// Parses the crop string (WIDTHxHEIGHT@X,Y)
void parse_crop(char *str)
{
	char *x = strchr(str, 'x');
	char *at = strchr(str, '@');
	char *comma = strchr(str, ',');

	if (!x || !at || !comma)
	{
		printf("Invalid crop format.\n");
		usage();
	}

	*x = *at = *comma = '\0';
	crop_w = atoi(str);
	crop_h = atoi(x+1);
	crop_x = atoi(at+1);
	crop_y = atoi(comma+1);
}

int main(int argc, char *argv[])
{

	if (argc != 15) usage();

	// Parse options
	parse_crop(argv[5]);
	const char *outfn = argv[6];
	tlx = atoi(argv[7]);
	tly = atoi(argv[8]);
	trx = atoi(argv[9]);
	try = atoi(argv[10]);
	blx = atoi(argv[11]);
	bly = atoi(argv[12]);
	brx = atoi(argv[13]);
	bry = atoi(argv[14]);

	// Open the files
	tiff_in_t tiffs;
	for (int col = 0; col < 4; ++col)
	{
		tiffs.chan[col] = tiff_open_r(argv[col+1]);
		if (!tiffs.chan[col])
		{
			die("Could not open channel: %s!\n", argv[col+1]);
		}
	}
	FILE *out_rgb = fopen(outfn, "w");
	if (!out_rgb) die("%s: Could not open fro writing.\n", outfn);

	// Do the work
	merge(&tiffs, out_rgb);

	// Close the files
	for (int col = 0; col < 4; ++col)
		tiff_close(tiffs.chan[col]);
	fclose(out_rgb);

	return 0;
}
